import os
from collections import namedtuple
from datetime import datetime

TrainRecord = namedtuple("TrainRecord", ["record_metric", "model_filename"])

class BestCheckpointSaver:
    def __init__(self, max_checkpoint_num=2, outputs="./outputs") -> None:
        self._outputs_dir = outputs
        self._max_checkpoint_num = max_checkpoint_num
        self._time_prefix = datetime.now().strftime(r"%m%d%H%M")
        
        self._best_checkpoint = []
        pass
    def update(self, trainer, metric_value):
        target_filename = f"{self._time_prefix}-{metric_value:.4f}.checkpoint"
        full_target_filepath = os.path.join(self._outputs_dir, target_filename)
        
        if len(self._best_checkpoint) <= self._max_checkpoint_num:
            trainer.save_checkpoint(full_target_filepath)
            self._best_checkpoint.append(TrainRecord(metric_value, full_target_filepath))
        else:
            self._best_checkpoint.sort(key=lambda x: x.record_metric, reverse=True)
            if metric_value > self._best_checkpoint[-1].record_metric:
                trainer.save_checkpoint(full_target_filepath)
                os.remove(self._best_checkpoint.pop().model_filename)
                self._best_checkpoint.append(TrainRecord(metric_value, full_target_filepath))
        
        return None